from omegaconf import DictConfig

from torch import nn

from ecgcmr.multimodal.multimodal_loss.GlobalMultiModalLoss import GlobalContrastiveLoss
from ecgcmr.multimodal.multimodal_loss.LocalMultiModalLossSoft import LocalContrastiveLossIn, LocalContrastiveLossOut, LocalLossCosineSimilarity


class MultiModalLoss(nn.Module):
    def __init__(self, cfg: DictConfig, loss_type: str, T_contrastive: int = 5) -> None:
        super().__init__()

        self.loss_type = loss_type
        self.weight_local = cfg.training_mode.loss.weight_local
        self.local_loss_type = cfg.training_mode.loss.local_loss.type

        assert self.loss_type in ["global", "local", "both"], "Invalid loss type"

        if self.loss_type in ["global", "both"]:
            self.global_loss = GlobalContrastiveLoss(
                temperature=cfg.training_mode.loss.global_loss.temperature
            )
        else:
            self.global_loss = None

        if self.loss_type in ["local", "both"]:
            if self.local_loss_type == "local_in":
                self.local_loss = LocalContrastiveLossIn(
                    T=T_contrastive,
                    temperature=cfg.training_mode.loss.local_loss.temperature,
                    std=cfg.training_mode.loss.local_loss.std,
                )
            elif self.local_loss_type == "local_out":
                self.local_loss = LocalContrastiveLossOut(
                    T=T_contrastive,
                    temperature=cfg.training_mode.loss.local_loss.temperature,
                    std=cfg.training_mode.loss.local_loss.std,
                )
            elif self.local_loss_type == "cosine_sim":
                self.local_loss = LocalLossCosineSimilarity()
            else:
                raise ValueError(f"Unsupported local_loss_type: {self.local_loss_type}")
        else:
            self.local_loss = None

    def forward(self, ecg_embeddings: dict, image_embeddings: dict):
        global_loss = None
        local_loss = None

        if self.global_loss is not None:
            global_loss = self.global_loss(
                ecg_embeddings["proj_ecg_global_emb"],
                image_embeddings["proj_img_global_emb"]
            )

        if self.local_loss is not None and self.weight_local > 0:
            local_loss = self.local_loss(
                ecg_embeddings["proj_ecg_local_emb"],
                image_embeddings["proj_img_local_emb"]
            )

        if self.loss_type == "global":
            total_loss = global_loss
        elif self.loss_type == "local":
            total_loss = local_loss
        else:
            total_loss = global_loss
            if local_loss is not None:
                total_loss = total_loss + self.weight_local * local_loss

        return {
            "global_loss": global_loss,
            "local_loss": local_loss,
            "total_loss": total_loss
        }
